package com.example.hook.monitor;

import javassist.*;

import javax.management.*;
import javax.management.remote.JMXConnectorServer;
import javax.management.remote.JMXConnectorServerFactory;
import javax.management.remote.JMXServiceURL;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.management.ManagementFactory;
import java.rmi.registry.LocateRegistry;
import java.security.ProtectionDomain;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 实时监控Agent - 可以动态启停监控
 */
public class MonitorAgent implements MonitorAgentMBean {
    
    private static Instrumentation instrumentation;
    private static MonitorTransformer transformer;
    private static boolean isMonitoring = false;
    private static MonitorAgent instance;
    
    // 监控统计
    private static final AtomicLong callCount = new AtomicLong(0);
    private static final AtomicLong totalExecutionTime = new AtomicLong(0);
    private static volatile boolean enableTimeTracking = true;
    
    public static void premain(String agentArgs, Instrumentation inst) {
        System.out.println("=== 监控Agent启动 ===");
        instrumentation = inst;
        instance = new MonitorAgent();
        
        // 注册MBean以便远程控制
        registerMBean();
        
        // 启动JMX服务器
        startJMXServer();
        
        // 默认开启监控
        if (instance != null) {
            instance.startMonitoring();
        }
        
        System.out.println("监控Agent初始化完成 - JMX端口: 9999");
        System.out.println("可以通过JConsole或JMX客户端连接: service:jmx:rmi:///jndi/rmi://localhost:9999/jmxrmi");
    }
    
    public static void agentmain(String agentArgs, Instrumentation inst) {
        premain(agentArgs, inst);
    }
    
    private static void registerMBean() {
        try {
            MBeanServer server = ManagementFactory.getPlatformMBeanServer();
            ObjectName name = new ObjectName("com.example.hook:type=MonitorAgent");
            server.registerMBean(instance, name);
            System.out.println("MBean注册成功: " + name);
        } catch (Exception e) {
            System.err.println("MBean注册失败: " + e.getMessage());
        }
    }
    
    private static void startJMXServer() {
        try {
            // 创建RMI注册表 - 使用不同的端口避免冲突
            LocateRegistry.createRegistry(9998);
            
            // 创建JMX连接器服务器
            MBeanServer mbs = ManagementFactory.getPlatformMBeanServer();
            JMXServiceURL url = new JMXServiceURL("service:jmx:rmi:///jndi/rmi://localhost:9998/jmxrmi");
            JMXConnectorServer cs = JMXConnectorServerFactory.newJMXConnectorServer(url, null, mbs);
            cs.start();
            
            System.out.println("JMX服务器启动成功，端口: 9998");
        } catch (Exception e) {
            System.err.println("JMX服务器启动失败: " + e.getMessage());
        }
    }
    
    @Override
    public void startMonitoring() {
        if (!isMonitoring && instrumentation != null) {
            transformer = new MonitorTransformer();
            instrumentation.addTransformer(transformer, true);
            
            // 重新转换已加载的类
            try {
                Class<?> targetClass = Class.forName("com.example.target.TargetService");
                instrumentation.retransformClasses(targetClass);
                isMonitoring = true;
                System.out.println("[Monitor] 监控已启动");
            } catch (Exception e) {
                System.err.println("[Monitor] 启动监控失败: " + e.getMessage());
            }
        }
    }
    
    @Override
    public void stopMonitoring() {
        if (isMonitoring && instrumentation != null && transformer != null) {
            instrumentation.removeTransformer(transformer);
            
            // 重新转换类以移除监控代码
            try {
                Class<?> targetClass = Class.forName("com.example.target.TargetService");
                instrumentation.retransformClasses(targetClass);
                isMonitoring = false;
                System.out.println("[Monitor] 监控已停止");
            } catch (Exception e) {
                System.err.println("[Monitor] 停止监控失败: " + e.getMessage());
            }
        }
    }
    
    @Override
    public boolean isMonitoring() {
        return isMonitoring;
    }
    
    @Override
    public long getCallCount() {
        return callCount.get();
    }
    
    @Override
    public long getAverageExecutionTime() {
        long count = callCount.get();
        return count > 0 ? totalExecutionTime.get() / count : 0;
    }
    
    @Override
    public void resetStatistics() {
        callCount.set(0);
        totalExecutionTime.set(0);
        System.out.println("[Monitor] 统计数据已重置");
    }
    
    @Override
    public void setTimeTrackingEnabled(boolean enabled) {
        enableTimeTracking = enabled;
        System.out.println("[Monitor] 时间追踪" + (enabled ? "已启用" : "已禁用"));
    }
    
    @Override
    public boolean isTimeTrackingEnabled() {
        return enableTimeTracking;
    }
    
    /**
     * 监控转换器
     */
    static class MonitorTransformer implements ClassFileTransformer {
        
        @Override
        public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                              ProtectionDomain protectionDomain, byte[] classfileBuffer) {
            
            if (className != null && className.equals("com/example/target/TargetService")) {
                if (isMonitoring) {
                    return addMonitoring(classfileBuffer);
                } else {
                    return removeMonitoring(classfileBuffer);
                }
            }
            
            return null;
        }
        
        private byte[] addMonitoring(byte[] classfileBuffer) {
            try {
                ClassPool pool = ClassPool.getDefault();
                CtClass cc = pool.makeClass(new java.io.ByteArrayInputStream(classfileBuffer));
                
                // 监控所有public方法
                CtMethod[] methods = cc.getDeclaredMethods();
                for (CtMethod method : methods) {
                    if (Modifier.isPublic(method.getModifiers())) {
                        addMonitoringToMethod(method);
                    }
                }
                
                byte[] byteCode = cc.toBytecode();
                cc.detach();
                
                System.out.println("[Monitor] 已添加监控到 TargetService");
                return byteCode;
                
            } catch (Exception e) {
                System.err.println("[Monitor] 添加监控失败: " + e.getMessage());
                return null;
            }
        }
        
        private byte[] removeMonitoring(byte[] classfileBuffer) {
            // 返回原始字节码（移除监控）
            return classfileBuffer;
        }
        
        private void addMonitoringToMethod(CtMethod method) throws Exception {
            String methodName = method.getName();
            
            // 添加方法调用前的监控代码
            method.insertBefore("{ " +
                "System.out.println(\"[Monitor] 方法调用: " + methodName + "() - 线程: \" + Thread.currentThread().getName() + \" - 时间: \" + new java.util.Date()); " +
                "com.example.hook.monitor.MonitorAgent.incrementCallCount(); " +
                "if (com.example.hook.monitor.MonitorAgent.isTimeTrackingEnabledStatic()) { " +
                "    com.example.hook.monitor.MonitorAgent.setMethodStartTime(System.nanoTime()); " +
                "}" +
                "}");
            
            // 添加方法调用后的监控代码
            method.insertAfter("{ " +
                "if (com.example.hook.monitor.MonitorAgent.isTimeTrackingEnabledStatic()) { " +
                "    long endTime = System.nanoTime(); " +
                "    long duration = endTime - com.example.hook.monitor.MonitorAgent.getMethodStartTime(); " +
                "    com.example.hook.monitor.MonitorAgent.addExecutionTime(duration); " +
                "    System.out.println(\"[Monitor] 方法执行完成: " + methodName + "() - 耗时: \" + (duration/1000000.0) + \"ms\"); " +
                "} else { " +
                "    System.out.println(\"[Monitor] 方法执行完成: " + methodName + "()\"); " +
                "}" +
                "}");
        }
    }
    
    // 静态方法供字节码注入使用
    private static final ThreadLocal<Long> methodStartTime = new ThreadLocal<>();
    
    public static void incrementCallCount() {
        callCount.incrementAndGet();
    }
    
    public static void setMethodStartTime(long time) {
        methodStartTime.set(time);
    }
    
    public static long getMethodStartTime() {
        Long time = methodStartTime.get();
        return time != null ? time : 0;
    }
    
    public static void addExecutionTime(long time) {
        totalExecutionTime.addAndGet(time);
    }
    
    public static boolean isTimeTrackingEnabledStatic() {
        return enableTimeTracking;
    }
}